from .network_architecture import *
from .loss import *

import os
import numpy as np
import tensorflow as tf

class Unet(object):
    """A Unet implementation"""

    def __init__(self, config):
        #self.network = mynetwork(config['network'])
        self.network = globals()[config['network']]

        tf.reset_default_graph()

        shape = config['input_shape']
        bs = config['batch_size']
        tensor_shape = np.hstack([bs, shape, 1])
        inputs_shape = np.hstack([bs, shape, 2])

        self.inputs = tf.placeholder("float", shape=inputs_shape)
        self.label = tf.placeholder("float", shape=tensor_shape)
        self.distance = tf.placeholder("float", shape=tensor_shape)
        self.iteration = tf.placeholder(tf.float32)
        self.keep_prob = tf.placeholder(tf.float32)
        self.training = tf.placeholder(tf.bool)

        output = self.network(self.inputs, self.training, self.keep_prob)
        self.cost = weighted_cost(self.label, output, self.distance)
        if config['weight_decay']:
            self.cost += tf.losses.get_regularization_loss()

        logits = tf.nn.sigmoid(output)
        self.prediction = tf.round(logits, name="prediction")

        # IOU
        intersection = tf.reduce_sum(self.prediction * self.label)
        union = tf.reduce_sum(self.prediction) + tf.reduce_sum(self.label)
        self.accuracy = 200 * intersection / (union + 1e-5)


    def predict(self, sess, inputs):
        prediction = sess.run(self.prediction,
                              feed_dict={self.inputs: inputs,
                                          self.training: False,
                                          self.keep_prob: 1.0, })
        return prediction

    def save(self, sess, save_path, epoch):
        """
        Save current model

        Args:
            sess: current session
            save_path: path to file system location
            epoch: number of epoch
        """

        iter_path = "{}/{}/epoch{}".format(save_path, self.network.__name__, str(int(epoch)))
        if not os.path.exists(iter_path):
            os.makedirs(iter_path)
        iter_path = os.path.join(iter_path, "model.cpkt")

        var_list = tf.trainable_variables()
        g_list = tf.global_variables()
        bn_moving_vars = [g for g in g_list if 'moving_mean' in g.name]
        bn_moving_vars += [g for g in g_list if 'moving_variance' in g.name]
        var_list += bn_moving_vars
        saver = tf.train.Saver(var_list=var_list, max_to_keep=5)

        save_path = saver.save(sess, iter_path)
        return save_path

    def _restore(self, sess, save_path):
        saver = tf.train.Saver()
        saver.restore(sess, save_path)
        print("Model restored from file: %s" % save_path)
